package cn.edu.szu;

import java.io.*;

public class Cipher {

    // 加密分组的大小，16个字节=128位
    private final int BLOCK_SIZE = 16;

    // 缓冲文件的字节大小，4096个字节=4kb
    private final int BUFFER_SIZE = 4096;

    // 指定原文件
    private final File originFile;

    // 加密后文件
    private final File encryptedFile;

    // 解密后文件
    private final File decryptFile;

    // 128位AES加密器
    private final AESUtils aesUtils;

    // 解密模式
    private final String mode;

    // 初始化向量
    private byte[] IV;

    // 构造函数
    public Cipher(String key, String mode, byte []IV, String filename) {
        // 创建加密器
        aesUtils = new AESUtils(key);

        // 初始化加密模式
        this.mode = mode;

        // 初始化向量
        this.IV = IV.clone();

        String filePath = "files\\";
        originFile = new File(filePath + filename);
        encryptedFile = new File(filePath + "E_" + filename);
        decryptFile = new File(filePath + "D_" + filename);

        // 新建文件
        creatNewFile(encryptedFile);
        creatNewFile(decryptFile);

    }

    // 新建文件
    private void creatNewFile(File file) {
        // 若文件已存在，新建文件
        if (file.exists()) {
            file.delete();
            try {
                file.createNewFile();
            } catch (IOException e) {
            }
        }
    }

    // 获取文件分组数量
    private int getBlockNumbers(File file) {
        return (int)(file.length() / BLOCK_SIZE);
    }

    // 文件加密
    public void fileEncrypt() {

        switch (mode) {
            case "ECB":
                ECB_Encrpyt();
                break;
            case "CBC":
                CBC_Encrypt();
                break;
            case "CFB":
                CFB_Encrypt();
                break;
            case "CTR":
                CTR_Xcrpyt(originFile, encryptedFile);
                break;
        }
    }

    // 文件解密
    public void fileDecrypt() {
        switch (mode) {
            case "ECB":
                ECB_Decrypt();
                break;
            case "CBC":
                CBC_Decrypt();
                break;
            case "CFB":
                CFB_Decrypt();
                break;
            case "CTR":
                CTR_Xcrpyt(encryptedFile, decryptFile);
            break;
        }
    }

    // 计算下一个计数器
    private void getNextCounter(byte[] counter) {
        int i = BLOCK_SIZE-1;
        if (counter[i] == 0xff) {
            counter[i] = 0;
            while (counter[i] == 0xff && i >=0) {
                counter[i] = 0;
                -- i;
            }
        } else {
            counter[i] = (byte) (counter[i] + 1);
        }
    }

    // CTR解密
    private void CTR_Xcrpyt(File in, File out) {
        // 初始向量作为计数器
        byte []counter = IV.clone();

        // 取得计数器的备份
        byte []copy_counter = counter.clone();

        // 密文
        byte []state = new byte[BLOCK_SIZE];

        // 字节缓冲区
        byte []buffer = new byte[BUFFER_SIZE];

        // 缓冲字节数
        int bufferSize = 0;

        // 每次读取的字节数
        int readLength = 0;

        try {
            InputStream input = new FileInputStream(in);

            // 当未读到文件末尾前加密
            while ((readLength = input.read(state)) != -1) {

                // 加密计数器并通过异或加密state
                aesUtils.encrypt(copy_counter);
                xorWithIv(state, copy_counter);

                // 使用缓冲区来提升文件读写效率
                bufferCopy(state, buffer, bufferSize);
                bufferSize += readLength;

                // 缓冲区满，写入文件同时字节数清零
                if (bufferSize == BUFFER_SIZE) {
                    write(out, buffer, bufferSize);
                    bufferSize = 0;
                }

                // 取得下个counter的备份
                getNextCounter(counter);
                copy_counter = counter.clone();
            }

            // 最后把剩余文件写入
            if (bufferSize != 0) {
                write(out, buffer, bufferSize);
            }

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    // CFB加密
    private void CFB_Encrypt() {
        // 明文或密文
        byte []state = new byte[BLOCK_SIZE];

        // 字节缓冲
        byte []buffer = new byte[BUFFER_SIZE];

        // 缓冲字节数
        int bufferSize = 0;

        // 读取字节数
        int readLength;

        // 将其初始化为初始化向量
        byte []pre_state = IV.clone();

        try {
            InputStream input = new FileInputStream(originFile);

            // 读到文件末尾前加密
            while ((readLength = input.read(state)) != -1) {

                // 加密前一个密文，第一次则加密IV
                aesUtils.encrypt(pre_state);

                // 执行异或操作来解密
                xorWithIv(state, pre_state);

                // 使用缓冲区来提升文件写入效率
                bufferCopy(state, buffer, bufferSize);
                bufferSize += readLength;

                // 若缓冲区满则写入文件
                if (bufferSize == BUFFER_SIZE) {
                    write(encryptedFile, buffer, bufferSize);
                    bufferSize = 0;
                }

                // 获得前一个密文
                pre_state = state.clone();
            }

            // 缓冲区未清空，则写入文件
            if (bufferSize != 0) {
                write(encryptedFile, buffer, bufferSize);
            }

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    // CFB解密
    private void CFB_Decrypt() {
        // 明文或密文
        byte []state = new byte[BLOCK_SIZE];

        // 字节缓冲
        byte []buffer = new byte[BUFFER_SIZE];

        // 缓冲字节数
        int bufferSize = 0;

        // 读取字节数
        int readLength;

        // 将其初始化为初始化向量
        byte []pre_state = IV.clone();

        // 用来拷贝前一组密文或明文
        byte []copy_pre_state;

        try {
            InputStream input = new FileInputStream(encryptedFile);

            // 未读到文件末尾时加密
            while ((readLength = input.read(state)) != -1) {

                // 获得当前密文的拷贝
                copy_pre_state = state.clone();

                // 加密前一个密文，第一次则加密IV
                aesUtils.encrypt(pre_state);

                // 执行异或操作来加密数据
                xorWithIv(state, pre_state);

                // 使用缓冲区来提升文件写入效率
                bufferCopy(state, buffer, bufferSize);
                bufferSize += readLength;

                // 若缓冲区满则写入文件
                if (bufferSize == BUFFER_SIZE) {
                    write(decryptFile, buffer, bufferSize);
                    bufferSize = 0;
                }

                // 获得前一个密文
                pre_state = copy_pre_state.clone();
            }

            // 缓冲区未清空，则写入文件
            if (bufferSize != 0) {
                write(decryptFile, buffer, bufferSize);
            }

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    // ECB加密
    private void ECB_Encrpyt() {

        // 明文
        byte []state = new byte[BLOCK_SIZE];

        // 字节缓冲
        byte []buffer = new byte[BUFFER_SIZE];

        // 缓冲字节数
        int bufferSize = 0;

        // 读取字节长度
        int readLength;

        try {
            InputStream input = new FileInputStream(originFile);

            // 未读到文件末尾前加密
            while ((readLength = input.read(state)) != -1) {
                if (readLength < BLOCK_SIZE) {
                    break;
                }

                // 加密明文
                aesUtils.encrypt(state);

                // 使用缓冲区提高效率
                bufferCopy(state, buffer, bufferSize);
                bufferSize += readLength;

                // 缓冲区满则写入加密文件
                if (bufferSize == BUFFER_SIZE) {
                    write(encryptedFile, buffer, bufferSize);
                    bufferSize = 0;
                }
            }

            // 若缓冲区不空，则写入加密文件
            if (bufferSize != 0) {
                write(encryptedFile, buffer, bufferSize);
            }

            // 处理文件末尾
            readLength = readLength % BLOCK_SIZE;
            // 文件填充
            padding(state, readLength);
            aesUtils.encrypt(state);
            write(encryptedFile, state, BLOCK_SIZE);

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    // ECB解密
    private void ECB_Decrypt() {

        // 密文
        byte []state = new byte[BLOCK_SIZE];

        // 字节缓冲
        byte []buffer = new byte[BUFFER_SIZE];

        // 缓冲字节数
        int bufferSize = 0;

        // 文件区块数量
        int blockNumbers = getBlockNumbers(encryptedFile);

        try {
            InputStream input = new FileInputStream(encryptedFile);

            // 解密区块，留下最后一个区块
            for (int i = 1; i < blockNumbers; ++i) {

                // 读取密文
                input.read(state);

                // 解密
                aesUtils.decrypt(state);

                // 写入缓冲区
                bufferCopy(state, buffer, bufferSize);
                bufferSize += BLOCK_SIZE;

                // 缓冲区满，写入文件
                if (bufferSize == BUFFER_SIZE) {
                    write(decryptFile, buffer, bufferSize);
                    bufferSize = 0;
                }
            }

            // 清空缓冲区
            if (bufferSize != 0) {
                write(decryptFile, buffer, bufferSize);
            }

            // 最后一个区块处理
            input.read(state);
            aesUtils.decrypt(state);
            // 反填充文件
            int lastLength = unPadding(state);
            write(decryptFile, state, lastLength);

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    // 异或操作
    private void xorWithIv(byte []state, byte []array) {
        for (int i=0; i<BLOCK_SIZE; ++i) {
            state[i] = (byte) (state[i] ^ array[i]);
        }
    }

    // CBC加密
    private void CBC_Encrypt() {

        // 明文
        byte []state = new byte[BLOCK_SIZE];

        // 字节缓冲
        byte []buffer = new byte[BUFFER_SIZE];

        // 缓冲字节数
        int bufferSize = 0;

        // 读取字节长度
        int readLength = 0;

        // 用于异或的向量
        byte []pre_state = IV.clone();

        try {
            InputStream input = new FileInputStream(originFile);

            // 未读到文件末尾前加密
            while ((readLength = input.read(state)) != -1) {
                // 与前一个密文异或，第一次与初始化向量异或
                xorWithIv(state, pre_state);
                if (readLength < BLOCK_SIZE) {
                    break;
                }

                // 加密明文
                aesUtils.encrypt(state);

                // 使用缓冲区提高效率
                bufferCopy(state, buffer, bufferSize);
                bufferSize += BLOCK_SIZE;

                // 缓冲区满则写入加密文件
                if (bufferSize == BUFFER_SIZE) {
                    write(encryptedFile, buffer, bufferSize);
                    bufferSize = 0;
                }

                // 获得前一个密文
                pre_state = state.clone();
            }

            // 若缓冲区不空，则写入加密文件
            if (bufferSize != 0) {
                write(encryptedFile, buffer, bufferSize);
            }

            // 计算最后一个区块长度
            readLength = readLength % BLOCK_SIZE;

            // 文件填充
            padding(state, readLength);
            xorWithIv(state, pre_state);

            // 加密
            aesUtils.encrypt(state);
            write(encryptedFile, state, BLOCK_SIZE);

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    // CBC解密
    private void CBC_Decrypt() {

        // 密文
        byte []state = new byte[BLOCK_SIZE];

        // 密文备份
        byte []copy_state;

        // 前一个密文
        byte []pre_state = IV.clone();

        // 字节缓冲
        byte []buffer = new byte[BUFFER_SIZE];

        // 缓存区字节
        int bufferSize = 0;

        // 文件区块数量
        int blockNumbers = getBlockNumbers(encryptedFile);

        try {
            InputStream input = new FileInputStream(encryptedFile);

            // 读到文件最后一个区块前
            for (int i = 1; i < blockNumbers; ++ i) {
                // 读取密文
                input.read(state);

                // 解密前要先保存前一轮的密文
                copy_state = state.clone();

                // 解密
                aesUtils.decrypt(state);

                // 异或
                xorWithIv(state, pre_state);

                // 将字节存入缓冲区以提高效率
                bufferCopy(state, buffer, bufferSize);
                bufferSize += BLOCK_SIZE;

                // 缓冲区满了，写入文件
                if (bufferSize == BUFFER_SIZE) {
                    write(decryptFile, buffer, bufferSize);
                    bufferSize = 0;
                }

                // 取出前一个密文
                pre_state = copy_state.clone();
            }

            // 清空缓冲区
            if (bufferSize != 0) {
                write(decryptFile, buffer, bufferSize);
            }

            // 读出最后一个区块
            input.read(state);
            // 解密
            aesUtils.decrypt(state);
            // 最后一轮异或
            xorWithIv(state, pre_state);
            // 反填充文件
            int lastLength = unPadding(state);
            write(decryptFile, state, lastLength);

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    // 赋值到一个字节缓冲数组，以提升读写文件效率
    private void bufferCopy(byte[] state, byte[] buffer, int start) {
        for (int i=0; i<BLOCK_SIZE; ++i) {
            buffer[start+i] = state[i];
        }
    }

    // 填充明文pkcs5算法
    private void padding(byte[] state, int length) {
        byte pad = (byte)(BLOCK_SIZE - length);
        for (int i=length; i<BLOCK_SIZE; ++i) {
            state[i] = pad;
        }
    }

    // 反填充明文pkcs5算法
    private int unPadding(byte[] state) {
        return BLOCK_SIZE - state[BLOCK_SIZE-1];
    }


    private void write(File file, byte[] state, int length) {
        OutputStream output = null;
        try {
            output = new FileOutputStream(file, true);
            output.write(state, 0, length);
            output.flush();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}
